/*
 *https://leetcode.cn/problems/count-the-number-of-good-nodes/description/
 *3249. 统计好节点的数目
 *medium 吴朝泽 2024.11.4
 *DFS深度优先搜索
*/

class Solution {
public:
    int countGoodNodes(vector<vector<int>>& edges) {
        unordered_map<int, vector<int>> tree;
        
        // 建立邻接表表示树
        for (auto& edge : edges) {
            tree[edge[0]].push_back(edge[1]);
            tree[edge[1]].push_back(edge[0]);
        }
        
        // 存储每个节点的子树大小
        vector<int> subtree_sizes(edges.size() + 1, 0);
        stack<pair<int, int>> stk;
        stk.push({0, -1});
        
        // 深度优先遍历，计算每个子树的大小
        while (!stk.empty()) {
            int node = stk.top().first, parent = stk.top().second;
            if (subtree_sizes[node] == 0) {  // 第一次访问该节点
                subtree_sizes[node] = 1;
                for (int child : tree[node]) {
                    if (child != parent) {
                        stk.push({child, node});
                    }
                }
            } else {  // 第二次访问该节点
                stk.pop();
                for (int child : tree[node]) {
                    if (child != parent) {
                        subtree_sizes[node] += subtree_sizes[child];  // 更新子树大小
                    }
                }
            }
        }
        
        int good_nodes_count = 0;
        stk.push({0, -1});
        
        // 再次遍历树，判断每个节点是否是"好节点"
        while (!stk.empty()) {
            int node = stk.top().first, parent = stk.top().second;
            stk.pop();
            if (tree[node].size() == 1 && node != 0) {  // 叶子节点
                good_nodes_count++;
                continue;
            }
            
            vector<int> sizes;
            for (int child : tree[node]) {
                if (child != parent) {
                    sizes.push_back(subtree_sizes[child]);
                }
            }
            
            // 判断该节点是否为"好节点"：所有子树大小是否相等
            bool allEqual = true;
            for (int i = 1; i < sizes.size(); ++i) {
                if (sizes[i] != sizes[0]) {
                    allEqual = false;
                    break;
                }
            }
            if (allEqual) {
                good_nodes_count++;  // 如果相等，则是"好节点"
            }
            
            // 继续遍历子节点
            for (int child : tree[node]) {
                if (child != parent) {
                    stk.push({child, node});
                }
            }
        }
        
        return good_nodes_count;  // 返回"好节点"的数量
    }
};
